# auth_system.py
import streamlit as st
import bcrypt
import json
import os
from pathlib import Path
import streamlit.components.v1 as components

# 文件路径设置
current_dir = Path(__file__).parent if "__file__" in locals() else Path.cwd()
users_file = current_dir / "users.json"

# 初始化用户数据
def init_users():
    if not users_file.exists():
        with open(users_file, "w") as f:
            json.dump({}, f)

# 加载用户数据
def load_users():
    with open(users_file) as f:
        return json.load(f)

# 保存用户数据
def save_users(users):
    with open(users_file, "w") as f:
        json.dump(users, f)

# 密码哈希处理
def hash_password(password):
    return bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()

# 验证密码
def verify_password(input_password, hashed_password):
    return bcrypt.checkpw(input_password.encode(), hashed_password.encode())

# 显示用户状态
def show_user_status():
    if st.session_state.get("logged_in"):
        user = st.session_state.current_user
        st.markdown(f"""
        <style>
            .user-info {{
                position: fixed;
                top: 10px;
                right: 10px;
                background: #f0f2f6;
                padding: 8px;
                border-radius: 5px;
                z-index: 999999;
                border: 1px solid #ddd;
            }}
            .user-info a {{
                color: #ff0000;
                text-decoration: none;
            }}
        </style>
        <div class="user-info">
            👤 {user['username']} ({user['role']}) 
            <a href="#" onclick="window.parent.document.querySelector('.stApp').dispatchEvent(new CustomEvent('logout'))">注销</a>
        </div>
        """, unsafe_allow_html=True)

# 注册页面
def register_page():
    st.header("新用户注册")
    with st.form("注册表单"):
        username = st.text_input("用户名")
        password = st.text_input("密码", type="password")
        confirm_password = st.text_input("确认密码", type="password")
        role = st.selectbox("用户角色", ["user", "admin"])
        
        if st.form_submit_button("注册"):
            users = load_users()
            if username in users:
                st.error("用户名已存在!")
                return
            if password != confirm_password:
                st.error("两次输入的密码不一致!")
                return
            if len(password) < 6:
                st.error("密码至少需要6个字符!")
                return
            
            users[username] = {
                "password": hash_password(password),
                "role": role
            }
            save_users(users)
            st.success("注册成功! 请登录")
            st.session_state.show_login = True

# 登录页面
def login_page():
    st.header("用户登录")
    with st.form("登录表单"):
        username = st.text_input("用户名")
        password = st.text_input("密码", type="password")
        
        if st.form_submit_button("登录"):
            users = load_users()
            user = users.get(username)
            
            if user and verify_password(password, user["password"]):
                st.session_state.logged_in = True
                st.session_state.current_user = {
                    "username": username,
                    "role": user["role"]
                }
                st.rerun()
            else:
                st.error("用户名或密码错误")

# 主内容页面
def main_content():
    st.success(f"欢迎回来, {st.session_state.current_user['username']}!")
    st.write("## 这是受保护的主页内容")
    
    if st.session_state.current_user["role"] == "admin":
        st.write("### 🛠️ 管理员专属功能")
        st.write("- 用户管理")
        st.write("- 系统设置")
        st.write("- 访问日志")
    else:
        st.write("### 👋 普通用户功能")
        st.write("- 查看个人资料")
        st.write("- 修改密码")

# 主程序
def main():
    init_users()
    st.set_page_config(
        page_title="用户认证系统",
        page_icon="🔒",
        layout="centered"
    )
    
    # 处理注销事件
    components.html(
        """
        <script>
        window.addEventListener('load', function() {
            window.parent.document.querySelector('.stApp').addEventListener('logout', function() {
                window.parent.document.querySelector('iframe').contentWindow.postMessage('logout', '*');
            });
        });
        </script>
        """, 
        height=0
    )
    
    if "logged_in" not in st.session_state:
        st.session_state.logged_in = False
    
    if st.session_state.logged_in:
        show_user_status()
        main_content()
    else:
        page = st.radio("请选择操作", ["登录", "注册"], horizontal=True)
        if page == "登录":
            login_page()
        else:
            register_page()

if __name__ == "__main__":
    main()